package cn.gzten.lambda.dynamo.annotation.processor;

import cn.gzten.annotation.processor.RegisterReflectionProcessor;
import cn.gzten.lambda.dynamo.annotation.DynamoBean;
import cn.gzten.lambda.dynamo.annotation.PartitionKey;
import com.google.auto.service.AutoService;
import com.squareup.javapoet.*;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import java.io.IOException;
import java.util.*;

@SupportedSourceVersion(SourceVersion.RELEASE_11)
@AutoService(Processor.class)
public class DynamoBeanProcessor extends AbstractProcessor {
    private Filer filer;
    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.filer = processingEnv.getFiler();
        System.out.println("DynamoBeanProcessor done initialization!");
    }

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Set.of(DynamoBean.class.getCanonicalName());
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        var elements = roundEnv.getElementsAnnotatedWith(DynamoBean.class);

        elements.forEach(e -> {
            DynamoBean an = e.getAnnotation(DynamoBean.class);

            try {
                String classFullName = e.toString();

                String keyName = null;
                for (Element el : e.getEnclosedElements()){
                    if (el instanceof VariableElement) {
                        var member = (VariableElement)el;
                        try {
                            PartitionKey annotationPartitionKey = member.getAnnotation(PartitionKey.class);
                            if (annotationPartitionKey != null) {
                                keyName = member.getSimpleName().toString();
                                break;
                            }
                        } catch (NullPointerException npe) {
                            continue;
                        }
                    }
                }
                if (keyName == null) {
                    throw new RuntimeException("Your class provided has no key field with @AppDynamoKey!");
                }

                classInfo classInfo = extractClassInfo(classFullName);

                FieldSpec fieldSpecForTableName = FieldSpec.builder(String.class, "TABLE_NAME")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                        .initializer("$S", an.tableName())
                        .build();

                FieldSpec fieldSpecForKeyName = FieldSpec.builder(String.class, "PARTITION_KEY")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                        .initializer("$S", keyName)
                        .build();

                FieldSpec fieldSpecForClass = FieldSpec.builder(Class.class, "CLASS")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                        .initializer(classFullName + ".class")
                        .build();
                MethodSpec getItem = MethodSpec.methodBuilder("getItem")
                        .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
                        .returns(TypeVariableName.get("java.util.Optional<" + classFullName + ">"))
                        .addParameter(software.amazon.awssdk.services.dynamodb.DynamoDbClient.class, "dynamo", Modifier.FINAL)
                        .addParameter(String.class, "partitionKey")
                        .addCode("var resp = dynamo.getItem(software.amazon.awssdk.services.dynamodb.model.GetItemRequest.builder()\n" +
                                "                    .tableName(TABLE_NAME)\n" +
                                "                    .key(cn.gzten.lambda.dynamo.util.AppTableMapper.singleMapWithString(PARTITION_KEY, partitionKey)).build());\n" +
                                "            if (resp.hasItem()) {\n" +
                                "                return cn.gzten.lambda.dynamo.util.AppTableMapper.itemToBean(resp.item(), CLASS);\n" +
                                "            } else {\n" +
                                "                return java.util.Optional.empty();\n" +
                                "            }")
                        .build();
                var spec = TypeSpec.classBuilder(classInfo.className + "DynamoBean")
                        .addOriginatingElement(e)
                        .addField(fieldSpecForTableName)
                        .addField(fieldSpecForKeyName)
                        .addField(fieldSpecForClass)
                        .addMethod(getItem)
                        .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                        .build();
                String packageName = classInfo.packageName;
                System.out.println(packageName);

                JavaFile.builder(packageName, spec).build().writeTo(filer);
            } catch (IOException ex) {
                ex.printStackTrace();
                throw new RuntimeException(ex);
            }

        });

        /**
         * Add Reflection for this also.
         */
        RegisterReflectionProcessor.processListOfClasses(elements, filer);
        return true;
    }

    public static final classInfo extractClassInfo(final String classFullName) {
        classInfo clz = new classInfo();
        if (classFullName.indexOf('.') > 0) {
            clz.packageName = classFullName.substring(0, classFullName.lastIndexOf('.'));
            clz.className = classFullName.substring(classFullName.lastIndexOf('.') + 1);
        } else {
            clz.packageName = "";
            clz.className = classFullName;
        }
        return clz;
    }

    private static class classInfo {
        public String packageName;
        public String className;
    }
}
